import warnings  # 用于控制警告的显示
import pandas as pd
from retrying import retry  # 用于实现重试机制，防止因偶发错误导致的操作失败
from dotenv import load_dotenv  # 用于加载环境变量，通常用于保护数据库连接的敏感信息
from sqlalchemy.dialects.mysql import insert  # sqlalchemy: 一个SQL工具包，支持连接各类数据库并执行SQL操作
from sqlalchemy import inspect, create_engine, text

load_dotenv()  # 加载环境变量
pd.set_option('display.max_rows', None)  # 显示所有行
pd.set_option('display.max_columns', None)  # 显示所有列
warnings.filterwarnings("ignore")  # 忽略警告
# DB_CONNECTION = os.getenv('DB_CONNECTION')
DB_NAME = "exhibition_data_test"  # 数据库名称
# 数据库连接信息
DB_CONNECTION = "mysql+pymysql://data_user:data_user#2024@mysql.xiangcunshenghuo.com:13306/exhibition_data_test"

_engine = None  # 使用全局变量_engine缓存数据库连接引擎，避免重复创建


def _get_engine():
    """ 获取连接对象 """
    try:
        global _engine
        if not _engine:
            _engine = create_engine(DB_CONNECTION, pool_size=10, max_overflow=50, pool_timeout=30, pool_recycle=-1)
        return _engine
    except Exception as e:
        raise Exception(f"_get_engine error: {e}")


def fetch_data(credit_code):
    """ 查询公司表中指定信用代码的记录 """
    try:
        engine = _get_engine()
        query = f"SELECT * FROM fh_company WHERE unified_social_credit_code = '{credit_code}'"
        df = pd.read_sql(query, con=engine)
        return df
    except Exception as e:
        raise Exception(f"fetch_data error: {e}")

def menu():
    """ 显示菜单 """
    print("1. 搜索")
    print("2. 批量搜索")
    print("3. 退出")
    choice = input("请输入选项：")
    if choice == "1":
        credit_code = input("请输入统一社会信用代码：")
        df = fetch_data(credit_code)
        print(df)
    elif choice == "2":
        print("批量搜索功能暂未实现")
    elif choice == "3":
        print("退出程序")
    else:
        print("输入错误，请重新输入")


if __name__ == '__main__':
    menu()
    # credit_code = '91330226MA2J7KC31U'
    # df = fetch_data(credit_code)
    # print(df)
